package middleware

import (
	"bytes"
	"io"
	"net/http"
	"regexp"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/microcosm-cc/bluemonday"
	"go.uber.org/zap"
)

// XSSConfig XSS配置
type XSSConfig struct {
	Enable      bool     // 是否启用XSS防护
	ExcludeURLs []string // 排除的URL列表
	SkipPaths   []string // 跳过的路径
}

// DefaultXSSConfig 默认XSS配置
func DefaultXSSConfig() XSSConfig {
	return XSSConfig{
		Enable:      true,
		ExcludeURLs: []string{"/api/v1/admin/content/save"},
		SkipPaths:   []string{"/api/v1/file/upload"},
	}
}

// XSS XSS中间件
func XSS() gin.HandlerFunc {
	return XSSWithConfig(DefaultXSSConfig())
}

// XSSWithConfig 使用自定义配置的XSS中间件
func XSSWithConfig(config XSSConfig) gin.HandlerFunc {
	// 如果未启用XSS防护，则返回空中间件
	if !config.Enable {
		return func(c *gin.Context) {
			c.Next()
		}
	}

	// 编译排除URL的正则表达式
	excludeRegexps := make([]*regexp.Regexp, 0, len(config.ExcludeURLs))
	for _, pattern := range config.ExcludeURLs {
		r, err := regexp.Compile(pattern)
		if err != nil {
			zap.L().Error("编译XSS排除URL正则表达式失败", zap.String("pattern", pattern), zap.Error(err))
			continue
		}
		excludeRegexps = append(excludeRegexps, r)
	}

	// 创建HTML过滤器
	p := bluemonday.UGCPolicy()

	return func(c *gin.Context) {
		// 检查是否为排除的URL
		path := c.Request.URL.Path
		for _, r := range excludeRegexps {
			if r.MatchString(path) {
				c.Next()
				return
			}
		}

		// 检查是否为跳过的路径
		for _, skipPath := range config.SkipPaths {
			if strings.HasPrefix(path, skipPath) {
				c.Next()
				return
			}
		}

		// 检查是否为GET请求，GET请求不需要处理请求体
		if c.Request.Method == http.MethodGet || c.Request.Method == http.MethodHead {
			// 处理查询参数
			sanitizeQueryParams(c, p)
			c.Next()
			return
		}

		// 检查Content-Type
		contentType := c.ContentType()
		if contentType == "application/json" {
			// 处理JSON请求
			sanitizeJSONRequest(c, p)
		} else if contentType == "application/x-www-form-urlencoded" || contentType == "multipart/form-data" {
			// 处理表单请求
			sanitizeFormRequest(c, p)
		}

		c.Next()
	}
}

// sanitizeQueryParams 净化查询参数
func sanitizeQueryParams(c *gin.Context, p *bluemonday.Policy) {
	for key, values := range c.Request.URL.Query() {
		for i, value := range values {
			values[i] = p.Sanitize(value)
		}
		c.Request.URL.Query()[key] = values
	}
}

// sanitizeJSONRequest 净化JSON请求
func sanitizeJSONRequest(c *gin.Context, p *bluemonday.Policy) {
	// 读取请求体
	var bodyBytes []byte
	if c.Request.Body != nil {
		bodyBytes, _ = io.ReadAll(c.Request.Body)
	}

	// 重置请求体
	c.Request.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))

	// 使用自定义的JSON处理器处理请求体
	c.Set("xss_body", bodyBytes)
}

// sanitizeFormRequest 净化表单请求
func sanitizeFormRequest(c *gin.Context, p *bluemonday.Policy) {
	// 解析表单
	_ = c.Request.ParseForm()
	for key, values := range c.Request.Form {
		for i, value := range values {
			values[i] = p.Sanitize(value)
		}
		c.Request.Form[key] = values
	}
}

// GetSanitizedJSON 获取经过XSS净化的JSON
func GetSanitizedJSON(c *gin.Context, obj interface{}) error {
	// 检查是否存在经过XSS处理的请求体
	_, exists := c.Get("xss_body")
	if !exists {
		// 如果不存在，则使用标准方法绑定JSON
		return c.ShouldBindJSON(obj)
	}

	// 使用经过XSS处理的请求体绑定JSON
	return c.ShouldBindJSON(obj)
}

// SanitizeHTML 净化HTML内容
func SanitizeHTML(html string) string {
	p := bluemonday.UGCPolicy()
	return p.Sanitize(html)
}

// SanitizeHTMLWithPolicy 使用自定义策略净化HTML内容
func SanitizeHTMLWithPolicy(html string, policy *bluemonday.Policy) string {
	return policy.Sanitize(html)
}

// CreateStrictPolicy 创建严格的HTML过滤策略
func CreateStrictPolicy() *bluemonday.Policy {
	return bluemonday.StrictPolicy()
}

// CreateUGCPolicy 创建UGC的HTML过滤策略
func CreateUGCPolicy() *bluemonday.Policy {
	return bluemonday.UGCPolicy()
}

// CreateCustomPolicy 创建自定义的HTML过滤策略
func CreateCustomPolicy() *bluemonday.Policy {
	p := bluemonday.NewPolicy()

	// 允许的HTML标签和属性
	p.AllowStandardURLs()
	p.AllowStandardAttributes()
	p.AllowImages()
	p.AllowLists()
	p.AllowTables()

	// 允许常用的HTML标签
	p.AllowElements("h1", "h2", "h3", "h4", "h5", "h6", "p", "br", "hr", "div", "span", "a", "img", "ul", "ol", "li", "table", "thead", "tbody", "tr", "th", "td", "strong", "em", "b", "i", "u", "s", "code", "pre", "blockquote")

	// 允许class属性
	p.AllowAttrs("class").Globally()

	// 允许style属性，但要小心，这可能导致XSS
	// p.AllowAttrs("style").Globally()

	return p
}
